package com.navinfo.tripanalysis.service.impl;

import com.navinfo.tripanalysis.pojo.Point;
import com.navinfo.tripanalysis.service.PointsSortGroupService;
import org.apache.spark.Partitioner;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFlatMapFunction;
import scala.Tuple2;

import java.io.Serializable;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/**
 * @author 沈东生
 */
public class PointsSortGroupServiceImpl2 implements PointsSortGroupService, Serializable {
    //并行度，重新分多少个区
    private Integer parallel;

    public Integer getParallel() {
        return parallel;
    }

    public void setParallel(Integer parallel) {
        this.parallel = parallel;
    }

    @Override
    public JavaPairRDD<Long, List<Point>> sortAndGroup(JavaRDD<Point> pointsRdd) {
        int partitionNum = Optional.ofNullable(getParallel()).orElse(1024);

        return pointsRdd.mapPartitionsToPair((PairFlatMapFunction<Iterator<Point>, Long, Point>) pointIterator ->
            new Iterator<Tuple2<Long, Point>>() {
                @Override
                public boolean hasNext() {
                    return pointIterator.hasNext();
                }
                @Override
                public Tuple2<Long, Point> next() {
                    Point p = pointIterator.next();
                    return new Tuple2<>(p.getTid(), p);
                }
            }).partitionBy(new TidPartitioner(partitionNum)).combineByKey(
                (Function<Point, List<Point>>) point -> Stream.of(point).collect(Collectors.toList()),
                (Function2<List<Point>, Point, List<Point>>) (points, p) -> {
                    points.add(p);
                    return points;
                },
                (Function2<List<Point>, List<Point>, List<Point>>) (p1, p2) -> {
                    p1.addAll(p2);
                    return p1;
                }).mapPartitionsToPair((PairFlatMapFunction<Iterator<Tuple2<Long, List<Point>>>, Long, List<Point>>) tuple2Iterator -> {
                    Iterable<Tuple2<Long, List<Point>>> iterator = () -> tuple2Iterator;
                    return StreamSupport.stream(iterator.spliterator(), false)
                            .map(t -> {
                                List<Point> collect = t._2.stream()
                                        .filter(dataInfo -> dataInfo.getGpsTime() > 0)
                                        .sorted(Point.COMPARATOR)
                                        .collect(Collectors.toList());
                                return new Tuple2<>(t._1, collect);
                            }).iterator();
                });
    }

    static class TidPartitioner extends Partitioner {
        private int partitionNum = 1 << 10;
        public TidPartitioner() {
        }
        public TidPartitioner(int partitionNum) {
            this.partitionNum = partitionNum;
        }
        @Override
        public int numPartitions() {
            return this.partitionNum;
        }
        @Override
        public int getPartition(Object key) {
            try {
                return (int) (Long.parseLong(key.toString()) % partitionNum);
            } catch (NumberFormatException e) {
                e.printStackTrace();
            }
            return 0;
        }
    }
}
